import torch
import torch.nn as nn


class MyLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        """
        这4个参数与传统RNN相同
        :param input_size: 输入的最后一个维度
        :param hidden_size: 隐藏层的最后一个维度
        :param output_size: 最后线性层的输出维度
        :param num_layers: 网络的层数
        """
        super(MyLSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers

        # 实例化
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
        # 实例化全连接线性层，用于将输出维度转化为指定的输出维度
        self.linear = nn.Linear(hidden_size, output_size)
        # 实例化nn中预定义的softmax层, 用于从输出层中获得类别的结果
        self.softmax = nn.LogSoftmax(-1)

    def forward(self, input1, hidden, c):
        """
        :param input1: 输入张量，形状：1 * n_letters
        :param hidden: 隐藏层张量，形状是 self.num_layers * 1 * self.hidden_size
        :param c: LSTM比传统RNN多了一个参数c，表示细胞状态张量
        :return:
        """
        input1 = input1.unsqueeze(0)
        # 1
        rr, (hn, c) = self.lstm(input1, (hidden, c))
        return self.softmax(self.linear(rr)), hn, c

    def init_hidden_c(self):
        # 初始化一个全0的隐藏层张量，维度是3
        c = hidden = torch.zeros(self.num_layers, 1, self.hidden_size)
        return hidden, c
